﻿import pandas as pd
import numpy as np
import pynimate as nim
import matplotlib.pyplot as plt
import matplotlib.dates as mdates
from datetime import datetime

# 创建模拟数据（实际应用中应使用真实历史数据）
def generate_population_data():
    # 创建年份范围（公元前3000年到公元2023年）
    years = list(range(-3000, 2024, 50))  # 每50年一个数据点
    
    # 主要文明/地区
    regions = [
        'Mesopotamia', 'Egypt', 'Indus Valley', 'China', 
        'Greece', 'Rome', 'Persia', 'India', 
        'Mesoamerica', 'Andes', 'Europe', 'Africa',
        'Islamic Caliphates', 'Mongol Empire', 'Ottoman Empire',
        'Global Population'
    ]
    
    # 创建空数据框
    df = pd.DataFrame(index=years, columns=regions)
    
    # 填充数据（模拟历史人口变化）
    for year in years:
        # 全球人口基线（基于历史估计）
        if year < -2000:
            global_pop = 10 + 0.0001 * (year + 3000)**2
        elif year < 0:
            global_pop = 30 + 0.0005 * (year + 2000)**2
        elif year < 1000:
            global_pop = 200 + 0.001 * year**2
        elif year < 1500:
            global_pop = 400 + 0.002 * (year - 1000)**2
        elif year < 1800:
            global_pop = 600 + 0.005 * (year - 1500)**2
        elif year < 1900:
            global_pop = 1000 + 0.01 * (year - 1800)**2
        elif year < 1950:
            global_pop = 2000 + 0.05 * (year - 1900)**2
        else:
            # 现代人口爆炸
            global_pop = 2500 + 2.5 * (year - 1950)**1.5
    
        # 设置全球人口
        df.loc[year, 'Global Population'] = global_pop * 1e6  # 转换为百万
        
        # 分配各区域人口（基于历史比例）
        if year < -2000:
            # 早期文明
            df.loc[year, 'Mesopotamia'] = global_pop * 0.15 * 1e6
            df.loc[year, 'Egypt'] = global_pop * 0.15 * 1e6
            df.loc[year, 'Indus Valley'] = global_pop * 0.15 * 1e6
            df.loc[year, 'China'] = global_pop * 0.15 * 1e6
            df.loc[year, 'Others'] = global_pop * 0.4 * 1e6
        elif year < 0:
            # 古典时期
            df.loc[year, 'Mesopotamia'] = global_pop * 0.1 * 1e6
            df.loc[year, 'Egypt'] = global_pop * 0.1 * 1e6
            df.loc[year, 'Persia'] = global_pop * 0.15 * 1e6
            df.loc[year, 'Greece'] = global_pop * 0.1 * 1e6
            df.loc[year, 'Rome'] = global_pop * 0.15 * 1e6
            df.loc[year, 'China'] = global_pop * 0.2 * 1e6
            df.loc[year, 'India'] = global_pop * 0.15 * 1e6
            df.loc[year, 'Others'] = global_pop * 0.05 * 1e6
        elif year < 1000:
            # 中世纪
            df.loc[year, 'China'] = global_pop * 0.25 * 1e6
            df.loc[year, 'India'] = global_pop * 0.2 * 1e6
            df.loc[year, 'Islamic Caliphates'] = global_pop * 0.2 * 1e6
            df.loc[year, 'Europe'] = global_pop * 0.15 * 1e6
            df.loc[year, 'Africa'] = global_pop * 0.1 * 1e6
            df.loc[year, 'Mesoamerica'] = global_pop * 0.05 * 1e6
            df.loc[year, 'Others'] = global_pop * 0.05 * 1e6
        elif year < 1500:
            # 中世纪晚期
            df.loc[year, 'China'] = global_pop * 0.3 * 1e6
            df.loc[year, 'India'] = global_pop * 0.2 * 1e6
            df.loc[year, 'Islamic Caliphates'] = global_pop * 0.15 * 1e6
            df.loc[year, 'Europe'] = global_pop * 0.15 * 1e6
            df.loc[year, 'Africa'] = global_pop * 0.1 * 1e6
            df.loc[year, 'Mesoamerica'] = global_pop * 0.05 * 1e6
            df.loc[year, 'Andes'] = global_pop * 0.03 * 1e6
            df.loc[year, 'Others'] = global_pop * 0.02 * 1e6
        elif year < 1800:
            # 近代早期
            df.loc[year, 'China'] = global_pop * 0.35 * 1e6
            df.loc[year, 'India'] = global_pop * 0.2 * 1e6
            df.loc[year, 'Ottoman Empire'] = global_pop * 0.1 * 1e6
            df.loc[year, 'Europe'] = global_pop * 0.15 * 1e6
            df.loc[year, 'Africa'] = global_pop * 0.1 * 1e6
            df.loc[year, 'Mesoamerica'] = global_pop * 0.02 * 1e6
            df.loc[year, 'Andes'] = global_pop * 0.02 * 1e6
            df.loc[year, 'Others'] = global_pop * 0.06 * 1e6
        else:
            # 现代
            df.loc[year, 'China'] = global_pop * 0.22 * 1e6
            df.loc[year, 'India'] = global_pop * 0.18 * 1e6
            df.loc[year, 'Europe'] = global_pop * 0.15 * 1e6
            df.loc[year, 'Africa'] = global_pop * 0.15 * 1e6
            df.loc[year, 'Others'] = global_pop * 0.3 * 1e6
    
    # 填充缺失值为0
    df = df.fillna(0)
    
    # 将年份转换为字符串（解决公元前年份问题）
    df.index = [f"{year} BC" if year < 0 else f"{year} AD" for year in df.index]
    
    return df

# 生成人口数据
df = generate_population_data()

# 创建画布 - 修正后的代码
cnv = nim.Canvas(
    figsize=(12, 8),  # 在这里设置画布大小
    facecolor='#f0f8ff'  # 淡蓝色背景
)

# 创建条形图竞赛
bar = nim.Barplot.from_df(
    df, 
    time_format='%Y', 
    ylabel='Population (Millions)',
    post_animate=False
)

# 自定义时间显示（处理公元前年份）
def format_year(i, datafier):
    return datafier.data.index[i]

bar.set_time(callback=format_year, fontsize=20, color='#333333')

# 设置标题和标签
bar.set_title("Human Population Growth Over 5000 Years", fontsize=24, weight='bold', color='#2c3e50')
bar.set_xlabel("Population (Millions)", fontsize=16)
bar.set_ylabel("Civilizations/Regions", fontsize=16)

# 自定义条形标签
bar.set_bar_annots(
    text_callback=lambda i, datafier: f"{datafier.data.iloc[i, 0]:.1f}M",
    fontsize=12,
    color='#333333'
)

# 设置颜色（为不同文明分配不同颜色）
colors = [
    '#1f77b4', '#ff7f0e', '#2ca02c', '#d62728', '#9467bd',
    '#8c564b', '#e377c2', '#7f7f7f', '#bcbd22', '#17becf',
    '#aec7e8', '#ffbb78', '#98df8a', '#ff9896', '#c5b0d5',
    '#c49c94'
]
bar.set_bar_color(colors)

# 添加图例
bar.set_legend(
    loc="upper left", 
    fontsize=10,
    labels=[
        'Mesopotamia', 'Egypt', 'Indus Valley', 'China', 
        'Greece', 'Rome', 'Persia', 'India', 
        'Mesoamerica', 'Andes', 'Europe', 'Africa',
        'Islamic Caliphates', 'Mongol Empire', 'Ottoman Empire',
        'Global Population'
    ]
)

# 添加历史事件注释
def add_historical_events(ax):
    """添加历史事件标记"""
    events = [
        (-3000, "First Civilizations", "Mesopotamia, Egypt, Indus Valley"),
        (-1200, "Bronze Age Collapse", "Major civilizations decline"),
        (-500, "Classical Age", "Greece, Persia, Rome"),
        (0, "Birth of Christ", "Start of Common Era"),
        (476, "Fall of Rome", "End of Western Roman Empire"),
        (622, "Rise of Islam", "Islamic Caliphates expand"),
        (1206, "Mongol Empire", "Largest contiguous empire"),
        (1492, "Columbus Voyage", "Age of Exploration begins"),
        (1750, "Industrial Revolution", "Rapid population growth"),
        (1914, "World War I", "Global conflict"),
        (1945, "World War II", "End of global conflict"),
        (1960, "Green Revolution", "Agricultural advancements"),
        (2020, "Modern Era", "Global population ~7.8 billion")
    ]
    
    for year, title, desc in events:
        # 转换年份为索引位置
        year_str = f"{year} BC" if year < 0 else f"{year} AD"
        
        # 如果年份在数据范围内
        if year_str in df.index:
            y_pos = df.loc[year_str].max() * 1.05
            ax.annotate(
                title,
                xy=(year_str, y_pos),
                xytext=(year_str, y_pos * 1.1),
                arrowprops=dict(facecolor='black', shrink=0.05),
                fontsize=10,
                ha='center'
            )
            ax.annotate(
                desc,
                xy=(year_str, y_pos * 0.95),
                fontsize=8,
                ha='center',
                color='#555555'
            )

# 添加数据来源注释
def add_data_source(ax):
    """添加数据来源说明"""
    ax.annotate(
        "Data: Historical estimates and projections\nSource: Various historical sources",
        xy=(0.5, 0.01),
        xycoords='figure fraction',
        ha='center',
        fontsize=9,
        color='#666666'
    )

# 添加动画回调函数
def post_animate_callback(ax):
    """动画后回调函数"""
    add_historical_events(ax)
    add_data_source(ax)
    
    # 添加网格
    ax.grid(True, linestyle='--', alpha=0.3)
    
    # 设置Y轴为对数刻度（更好地展示人口变化）
    ax.set_yscale('log')
    
    # 设置Y轴标签
    ax.set_yticks([1e6, 1e7, 1e8, 1e9])
    ax.set_yticklabels(['1M', '10M', '100M', '1B'])
    
    # 添加重要里程碑线
    milestones = {
        1e6: "1 Million",
        1e7: "10 Million",
        1e8: "100 Million",
        1e9: "1 Billion"
    }
    
    for value, label in milestones.items():
        ax.axhline(y=value, color='#cccccc', linestyle='--', alpha=0.5)
        ax.annotate(
            label,
            xy=(df.index[0], value),
            xytext=(-50, 0),
            textcoords='offset points',
            fontsize=8,
            color='#777777',
            va='center'
        )

# 设置动画后回调
bar.post_animate = post_animate_callback

# 添加图表到画布
cnv.add_plot(bar)

# 添加时间滑块
slider = nim.TimeSlider(bar, color='#3498db')
cnv.add_plot(slider)

# 添加总人口折线图
def add_global_population_line(ax):
    """添加全球人口折线图"""
    # 获取全球人口数据
    global_pop = df['Global Population']
    
    # 创建次坐标轴
    ax2 = ax.twinx()
    
    # 绘制折线图
    line, = ax2.plot(
        global_pop.index, 
        global_pop.values, 
        color='#e74c3c', 
        linewidth=2.5,
        label='Global Population'
    )
    
    # 设置次坐标轴属性
    ax2.set_yscale('log')
    ax2.set_ylim(1e6, 1e10)
    ax2.set_yticks([1e6, 1e7, 1e8, 1e9, 1e10])
    ax2.set_yticklabels(['1M', '10M', '100M', '1B', '10B'])
    ax2.set_ylabel('Global Population', fontsize=12, color='#e74c3c')
    ax2.tick_params(axis='y', colors='#e74c3c')
    
    # 添加图例
    ax2.legend(
        handles=[line], 
        loc='upper right',
        fontsize=10
    )
    
    return ax2

# 添加折线图回调
bar.post_animate = lambda ax: (
    post_animate_callback(ax),
    add_global_population_line(ax)
)

# 生成动画 - 在这里设置动画参数
cnv.animate(
    interval=100,       # 帧间隔（毫秒）
    blit=True,          # 使用blitting优化性能
    repeat=True,        # 动画循环播放
    repeat_delay=3000   # 循环之间的延迟
)

# 保存动画
cnv.save("human_population_5000_years.gif", fps=15, dpi=150)

# 显示动画（在Jupyter中）
# cnv.show()